from langchain.schema.runnable import Runnable
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableMap
from langserve import RemoteRunnable

langserve_chat_api = RemoteRunnable("http://localhost:8000/chat/")
prompt = ChatPromptTemplate.from_messages(
    [("system", "你是一个喜欢写故事的助手"), ("user", "{input}")]
)

# 自定义链
chain = prompt | RunnableMap({
    "chat_api": langserve_chat_api
})

print("同步调用langserve /chat结果")
response = chain.invoke({"input":"写一个关于护士的故事"})
print(response)
